{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import sys\n",
    "sys.path.insert(0,'../')\n",
    "from pose_graph_partitioning.pose_graph import *\n",
    "from pose_graph_partitioning.pose_graph_partitioning import *\n",
    "from simulate_utils import ATE, align_posegraphs, call_d2pgo_opti\n",
    "from subprocess import PIPE, Popen\n",
    "import copy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def call_d2pgo(input_data, output_data, agent_num, params, param_fixed={}, cmd=\"roslaunch d2pgo d2pgo_test_multi.launch \"):\n",
    "    command = f\"{cmd} g2o_path:={input_data} output_path:={output_data} agent_num:={agent_num} \"\n",
    "    for k, v in params:\n",
    "        command += f\"{k}:={v[0]} \"\n",
    "    for k in param_fixed:\n",
    "        v = param_fixed[k]\n",
    "        command += f\"{k}:={v} \"\n",
    "    p = Popen(command, shell=True, stdout=PIPE, stderr=PIPE)\n",
    "    print(command)\n",
    "    stdout, stderr = p.communicate()\n",
    "    # return stdout, stderr\n",
    "    # print(stdout.decode(\"utf-8\"))\n",
    "    return\n",
    "\n",
    "def check_param_unique(params):\n",
    "    cnt = 0\n",
    "    for k, v in params:\n",
    "        if len(v) > 1:\n",
    "            return False, cnt\n",
    "        cnt += 1\n",
    "    return True, -1 \n",
    "\n",
    "def print_res_to_table(res, sort_index=\"final_cost\"):\n",
    "    import tabulate\n",
    "    output_table = []\n",
    "    #Add title\n",
    "    if len(res) == 0:\n",
    "        return []\n",
    "    sort_values = []\n",
    "    for row in res:\n",
    "        output_table.append([])\n",
    "        for k in row:\n",
    "            v = row[k]\n",
    "            if isinstance(v, np.ndarray) or isinstance(v, list):\n",
    "                output_table[-1].append(v[0])\n",
    "            else:\n",
    "                if k == \"ATE_T\":\n",
    "                    output_table[-1].append(f\"{v:.3f}m\")\n",
    "                elif k == \"ATE_rot\":\n",
    "                    output_table[-1].append(f\"{v*57.3:.3f}\")\n",
    "                elif k == \"final_cost\":\n",
    "                    output_table[-1].append(f\"{v:.3f}\")\n",
    "                else:\n",
    "                    output_table[-1].append(v)\n",
    "            if k == sort_index:\n",
    "                sort_values.append(v)\n",
    "    output_table = np.array(output_table)\n",
    "    sort_values = np.array(sort_values)\n",
    "    output_table = output_table[sort_values.argsort()]\n",
    "    row0 = []\n",
    "    for key in res[0]:\n",
    "        row0.append(key)\n",
    "    output_table = output_table.tolist()\n",
    "    output_table.insert(0, row0)\n",
    "    return tabulate.tabulate(output_table, tablefmt='html')\n",
    "\n",
    "def grid_search(input_data, output_data, gt_data, agent_num, params, params_fixed={}, cmd=\"roslaunch d2pgo d2pgo_test_multi.launch \", is_root=True):\n",
    "    #Gird search in param lists\n",
    "    pg_gt = PoseGraph()\n",
    "    pg_gt.read_g2o_single(gt_data)\n",
    "    params_num = len(params)\n",
    "    #DFS on params\n",
    "    unique, row = check_param_unique(params)\n",
    "    if unique:\n",
    "        #Only one param, no need to grid search\n",
    "        call_d2pgo(input_data, output_data, agent_num, params, params_fixed, cmd)\n",
    "        pg_output = PoseGraph()\n",
    "        pg_output.read_g2o_folder(output_data, prt=False)\n",
    "        ate_T, ate_rot = ATE(pg_output, pg_gt)\n",
    "        # align_posegraphs(pg_output, pg_gt)\n",
    "        # ax = pg_gt.show(\"GT\")\n",
    "        # pg_output.show(\"output\", ax=ax, clear=False)\n",
    "        # print(params, ate_T, ate_rot)\n",
    "        ret_data = {}\n",
    "        for k, v in params:\n",
    "            ret_data[k] = v[0]\n",
    "        ret_data[\"ATE_T\"] = ate_T\n",
    "        ret_data[\"ATE_rot\"] = ate_rot\n",
    "        ret_data[\"final_cost\"] = pg_output.evaluate_cost()\n",
    "        # Add this data to output_data/result.txt\n",
    "        with open(output_data + \"/search_result.txt\", \"a\") as f:\n",
    "            # Save the params to a line of dict\n",
    "            f.write(str(ret_data) + \"\\n\")\n",
    "        return [ret_data]\n",
    "    else:\n",
    "        #More than one param, need to grid search\n",
    "        param_values = params[row][1]\n",
    "        res = []\n",
    "        for i in range(len(param_values)):\n",
    "            new_params = copy.deepcopy(params)\n",
    "            new_params[row][1] = [param_values[i]]\n",
    "            ret = grid_search(input_data, output_data, gt_data, agent_num, new_params, params_fixed, cmd, is_root=False)\n",
    "            for obj in ret:\n",
    "                res.append(obj)\n",
    "    return res\n",
    "\n",
    "def random_search(input_data, output_data, agent_num, params, gt_data=None, max_trials=100, params_fixed={}, cmd=\"roslaunch d2pgo d2pgo_test_multi.launch \", is_root=True):\n",
    "    #Gird search in param lists\n",
    "    if gt_data is not None:\n",
    "        pg_gt = PoseGraph()\n",
    "        pg_gt.read_g2o_single(gt_data)\n",
    "    #DFS on params\n",
    "    res = []\n",
    "    for i in range(max_trials):\n",
    "        # Take a random sample of params\n",
    "        ret_data = {}\n",
    "        params_instance = []\n",
    "        for param in params:\n",
    "            params_instance.append([param[0], [np.random.choice(param[1])]])\n",
    "            ret_data[param[0]] = params_instance[-1][1]\n",
    "        print(f\"trial {i}:\", params_instance)\n",
    "        #Only one param, no need to grid search\n",
    "        call_d2pgo(input_data, output_data, agent_num, params_instance, params_fixed, cmd)\n",
    "        pg_output = PoseGraph()\n",
    "        pg_output.read_g2o_folder(output_data, prt=False)\n",
    "        if gt_data is not None:\n",
    "            ate_T, ate_rot = ATE(pg_output, pg_gt)\n",
    "            ret_data[\"ATE_T\"] = ate_T\n",
    "            ret_data[\"ATE_rot\"] = ate_rot\n",
    "        ret_data[\"final_cost\"] = pg_output.evaluate_cost()\n",
    "        # Add this data to output_data/result.txt\n",
    "        with open(output_data + \"/search_result.txt\", \"a\") as f:\n",
    "            # Save the params to a line of dict\n",
    "            f.write(str(ret_data) + \"\\n\")\n",
    "        res.append(ret_data)\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_data = \"/home/xuhao/data/d2slam/pgo/parking-garage/input\"\n",
    "gt_data = \"/home/xuhao/data/d2slam/pgo/parking-garage/groundtruth.g2o\"\n",
    "output_data = \"/home/xuhao/data/d2slam/pgo/parking-garage/searching-output/\"\n",
    "agent_num = 5\n",
    "params = [\n",
    "    [\"max_steps\", [100]],\n",
    "    [\"rho_frame_T\", np.linspace(0.01, 2.0, 10000)],\n",
    "    [\"rho_frame_theta\", np.linspace(0.01, 2.0, 10000)],\n",
    "    [\"eta_k\", np.linspace(0.5, 3, 100)],\n",
    "    [\"enable_rot_init\", [\"True\", \"False\"]]\n",
    "]\n",
    "ret = random_search(input_data, output_data, agent_num, params, max_trials=1000)\n",
    "print_res_to_table(ret)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_data = \"/home/xuhao/data/d2slam/pgo/tum_corr_5/input/\"\n",
    "gt_data = \"/home/xuhao/data/d2slam/pgo/tum_corr_5/input/fullGraph_optimized.g2o\"\n",
    "output_data = \"/home/xuhao/data/d2slam/pgo/tum_corr_5/searching-output/\"\n",
    "agent_num = 5\n",
    "# params = [\n",
    "#     [\"max_steps\", [100]],\n",
    "#     [\"rho_frame_T\", np.logspace(-5, -1, 10000, base=10)],\n",
    "#     [\"rho_frame_theta\", np.logspace(-2, 2, 10000, base=10)],\n",
    "#     [\"eta_k\", np.logspace(-1, 0, 100, base=10)],\n",
    "#     [\"enable_rot_init\", [\"false\", \"true\"]]\n",
    "# ]\n",
    "params = [\n",
    "    [\"max_steps\", [10000]],\n",
    "    [\"rho_frame_T\", np.linspace(0.01, 5.0, 10000)],\n",
    "    [\"rho_frame_theta\", np.linspace(0.01, 5.0, 10000)],\n",
    "    [\"eta_k\", np.linspace(0.5, 3, 100)],\n",
    "    [\"rho_rot_mat\", np.linspace(0.001, 3, 100)],\n",
    "    [\"enable_rot_init\", [\"True\"]]\n",
    "]\n",
    "ret = random_search(input_data, output_data, agent_num, params, max_trials=1000)\n",
    "print_res_to_table(ret)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Show the convergence of the results\n",
    "intput_data = \"/home/xuhao/data/d2slam/pgo/torus3D/input/\"\n",
    "gt_data = \"/home/xuhao/data/d2slam/pgo/torus3D/groundtruth.g2o\"\n",
    "output_data = \"/home/xuhao/data/d2slam/pgo/torus3D/output/\"\n",
    "agent_num = 5\n",
    "params = [\n",
    "    [\"max_steps\", np.logspace(0, 2, 10, base=10, dtype=int)],\n",
    "    [\"debug_rot_init_only\", [\"true\"]]\n",
    "]\n",
    "fixed_param = {\n",
    "    \"debug_rot_init_only\":\"true\"\n",
    "}\n",
    "ret = grid_search(intput_data, output_data, gt_data, agent_num, params, fixed_param)\n",
    "print_res_to_table(ret, \"max_steps\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "intput_data = \"/home/xuhao/data/d2slam/pgo/torus3D/input/\"\n",
    "gt_data = \"/home/xuhao/data/d2slam/pgo/torus3D/groundtruth.g2o\"\n",
    "output_data = \"/home/xuhao/data/d2slam/pgo/torus3D/output/\"\n",
    "agent_num = 5\n",
    "params = [\n",
    "    [\"rho_rot_mat\", np.logspace(-3, 1, 10, base=10)]\n",
    "]\n",
    "fixed_param = {\n",
    "    \"debug_rot_init_only\":\"true\",\n",
    "    \"max_steps\": 20\n",
    "}\n",
    "ret = grid_search(intput_data, output_data, gt_data, agent_num, params, fixed_param)\n",
    "print_res_to_table(ret, \"ATE_rot\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_function(param):\n",
    "    pg_output, _, _ = call_d2pgo_opti(input_data, output_data, agent_num = agent_num, max_steps=100, enable_rot_init=True, \n",
    "        eta_k=param['eta_k'], rho_frame_theta=param['rho_frame_theta'], rho_frame_T=param['rho_frame_T'])\n",
    "    return pg_output.evaluate_cost()\n",
    "# input_data = \"/home/xuhao/data/d2slam/pgo/tum_corr_5/input\"\n",
    "# output_data = \"/home/xuhao/data/d2slam/pgo/tum_corr_5/searching-output/\"\n",
    "\n",
    "input_data = \"/home/xuhao/data/d2slam/pgo/parking-garage/input\"\n",
    "output_data = \"/home/xuhao/data/d2slam/pgo/parking-garage/searching-output/\"\n",
    "\n",
    "agent_num = 5\n",
    "from hyperopt import hp, fmin, tpe\n",
    "space = {\n",
    "    \"rho_frame_T\": hp.uniform(\"rho_frame_T\", 0.0, 10.0),\n",
    "    \"rho_frame_theta\": hp.uniform(\"rho_frame_theta\", 0.0, 10.0),\n",
    "    \"eta_k\": hp.uniform(\"eta_k\", 0, 5.0),\n",
    "}\n",
    "\n",
    "best = fmin(search_function, space, algo=tpe.suggest, max_evals=100)\n",
    "print(best)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.12 ('base')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "08ce52785f0fedc81003ce387e097a83d6cc9494681cd746006386992005bb71"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
